package top.faig.utils;

import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.util.encoders.Hex;

import java.security.SecureRandom;
import java.security.Security;
import java.util.Arrays;

public class SM4Util {

  static {
    Security.addProvider(new BouncyCastleProvider());
  }

  private static final int BLOCK_SIZE = 16;

  /**
   * 生成一个随机的 SM4 密钥
   *
   * @return 生成长度 32 的密钥，返回的密钥以十六进制字符串形式表示
   */
  public static String generateKey() {
    byte[] key = new byte[BLOCK_SIZE];
    SecureRandom random = new SecureRandom();
    random.nextBytes(key);
    return Hex.toHexString(key);
  }

  /**
   * SM4加密
   *
   * @param data 需要加密的数据
   * @param key  密钥 32 位十六进制字符串表示的（16 字节）
   *
   * @return 加密后的数据
   * @throws InvalidCipherTextException 加密失败时抛出异常
   */
  public static String encrypt(String data, String key) throws InvalidCipherTextException {
    return encrypt(data, key, null);
  }

  /**
   * SM4解密
   *
   * @param encryptedData 加密的数据
   * @param key  密钥 32 位十六进制字符串表示的（16 字节）
   * @return 解密后的数据
   * @throws InvalidCipherTextException 解密失败时抛出异常
   */
  public static String decrypt(String encryptedData, String key) throws InvalidCipherTextException {
    return decrypt(encryptedData, key, null);
  }

  /**
   * SM4加密
   *
   * @param data 需要加密的数据
   * @param key 密钥 32 位十六进制字符串表示的（16 字节）
   * @param iv  初始化向量 32 位十六进制字符串表示的（16 字节）或为空，若为空则使用全零 IV
   *
   * @return 加密后的数据
   * @throws InvalidCipherTextException 加密失败时抛出异常
   */
  public static String encrypt(String data, String key, String iv) throws InvalidCipherTextException {
    byte[] keyBytes = Hex.decodeStrict(key);
    if (keyBytes.length != BLOCK_SIZE) {
      throw new IllegalArgumentException("Key 长度必须为 32");
    }

    byte[] ivBytes;
    if (iv == null || iv.isEmpty()) {
      ivBytes = new byte[BLOCK_SIZE]; // 默认使用全零 IV
    } else {
      ivBytes = Hex.decodeStrict(iv);
      if (ivBytes.length != BLOCK_SIZE) {
        throw new IllegalArgumentException("IV 如果不为空，长度必须为 32");
      }
    }

    SM4Engine sm4Engine = new SM4Engine();
    CBCBlockCipher cbcBlockCipher = new CBCBlockCipher(sm4Engine);
    PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(cbcBlockCipher);

    CipherParameters params = new ParametersWithIV(new KeyParameter(keyBytes), ivBytes);
    cipher.init(true, params);
    byte[] dataBytes = data.getBytes();
    byte[] process = process(cipher, dataBytes);
    return Hex.toHexString(process);
  }

  /**
   * SM4解密
   *
   * @param encryptedData 加密的数据
   * @param key 密钥 32 位十六进制字符串表示的（16 字节）
   * @param iv  初始化向量 32 位十六进制字符串表示的（16 字节）或为空，若为空则使用全零 IV
   * @return 解密后的数据
   * @throws InvalidCipherTextException 加密失败时抛出异常
   */
  public static String decrypt(String encryptedData, String key, String iv) throws InvalidCipherTextException {
    byte[] keyBytes = Hex.decodeStrict(key);
    if (keyBytes.length != BLOCK_SIZE) {
      throw new IllegalArgumentException("Key 长度必须为 32");
    }

    byte[] ivBytes;
    if (iv == null || iv.isEmpty()) {
      ivBytes = new byte[BLOCK_SIZE]; // 默认使用全零 IV
    } else {
      ivBytes = Hex.decodeStrict(iv);
      if (ivBytes.length != BLOCK_SIZE) {
        throw new IllegalArgumentException("IV 如果不为空，长度必须为 32");
      }
    }

    SM4Engine sm4Engine = new SM4Engine();
    CBCBlockCipher cbcBlockCipher = new CBCBlockCipher(sm4Engine);
    PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(cbcBlockCipher);

    CipherParameters params = new ParametersWithIV(new KeyParameter(keyBytes), ivBytes);
    cipher.init(false, params);
    byte[] encryptedDataBytes = Hex.decodeStrict(encryptedData);
    byte[] process = process(cipher, encryptedDataBytes);
    return new String(process);
  }

  /**
   * 处理加密或解密的数据
   *
   * @param cipher 加密器或解密器
   * @param data 数据
   *
   * @return 处理后的数据
   * @throws InvalidCipherTextException 处理失败时抛出异常
   */
  private static byte[] process(PaddedBufferedBlockCipher cipher, byte[] data) throws InvalidCipherTextException {
    byte[] outputBuffer = new byte[cipher.getOutputSize(data.length)];
    int length = cipher.processBytes(data, 0, data.length, outputBuffer, 0);
    length += cipher.doFinal(outputBuffer, length);
    return Arrays.copyOf(outputBuffer, length);
  }



  public static void main(String[] args) throws Exception {
    String key = generateKey();
    System.out.println("SM4 随机密钥" + key);

    String data = "Hello SM";
    String encrypt = encrypt(data, key);
    System.out.println("SM4 加密" + encrypt);
    String decryptData = decrypt(encrypt, key);
    System.out.println("SM4 解密" + decryptData);
  }

}
